package main

import (
	"fmt"
	"sync"
)

func producer(nums ...int) <-chan int {
	out := make(chan int)
	go func() {
		defer close(out)
		for i := range nums {
			out <- nums[i]
		}
	}()
	return out
}

func square(in <-chan int) <-chan int {
	out := make(chan int)
	go func() {
		defer close(out)
		for v := range in {
			out <- v * v
		}
	}()
	return out
}

func merge(ins ...<-chan int) <-chan int {
	out := make(chan int)
	wg := sync.WaitGroup{}
	wg.Add(len(ins))
	for i := range ins {
		go func(in <-chan int) {
			defer wg.Done()
			for v := range in {
				out <- v
			}
		}(ins[i])
	}

	go func() {
		wg.Wait()
		close(out)
	}()
	return out
}

func main() {
	in1 := producer(1, 2, 3)
	in2 := producer(4, 5, 6)

	in1out1 := square(in1)
	in1out2 := square(in1)

	in2out1 := square(in2)
	in2out2 := square(in2)

	for v := range merge(
		in1out1,
		in1out2,
		in2out1,
		in2out2) {
		fmt.Printf("%3d\n", v)
	}
}
